#!/usr/bin/env python3
"""
MERIT-Hydro 栅格数据下载脚本
从 https://mghydro.com/watersheds/rasters/ 下载流域栅格数据

使用方法:
    python merit_hydro_rasters.py
    nohup python merit_hydro_rasters.py > ../logs/rasters.log 2>&1 &
"""

import os
import sys
import requests
import concurrent.futures
from loguru import logger
from tqdm import tqdm
import threading

# 配置
BASE_URL = "https://mghydro.com/watersheds/rasters/"
# 使用相对于项目根目录的data文件夹
BASE_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "mghydro", "watersheds", "rasters")
MAX_WORKERS = 8

# 文件列表 - 直接从网站获取
FILES_TO_DOWNLOAD = {
    "accum_basins/": [
        "accum11.tif", "accum12.tif", "accum13.tif", "accum14.tif", "accum15.tif",
        "accum16.tif", "accum17.tif", "accum18.tif", "accum21.tif", "accum22.tif",
        "accum23.tif", "accum24.tif", "accum25.tif", "accum26.tif", "accum27.tif",
        "accum28.tif", "accum29.tif", "accum31.tif", "accum32.tif", "accum33.tif",
        "accum34.tif", "accum35.tif", "accum36.tif", "accum41.tif", "accum42.tif",
        "accum43.tif", "accum44.tif", "accum45.tif", "accum46.tif", "accum47.tif",
        "accum48.tif", "accum49.tif", "accum51.tif", "accum52.tif", "accum53.tif",
        "accum54.tif", "accum55.tif", "accum56.tif", "accum57.tif", "accum61.tif",
        "accum62.tif", "accum63.tif", "accum64.tif", "accum65.tif", "accum66.tif",
        "accum67.tif", "accum71.tif", "accum72.tif", "accum73.tif", "accum74.tif",
        "accum75.tif", "accum76.tif", "accum77.tif", "accum78.tif", "accum81.tif",
        "accum82.tif", "accum83.tif", "accum84.tif", "accum85.tif", "accum86.tif",
        "accum91.tif"
    ],
    "flow_dir_basins/": [
        "flowdir11.tif", "flowdir12.tif", "flowdir13.tif", "flowdir14.tif", "flowdir15.tif",
        "flowdir16.tif", "flowdir17.tif", "flowdir18.tif", "flowdir21.tif", "flowdir22.tif",
        "flowdir23.tif", "flowdir24.tif", "flowdir25.tif", "flowdir26.tif", "flowdir27.tif",
        "flowdir28.tif", "flowdir29.tif", "flowdir31.tif", "flowdir32.tif", "flowdir33.tif",
        "flowdir34.tif", "flowdir35.tif", "flowdir36.tif", "flowdir41.tif", "flowdir42.tif",
        "flowdir43.tif", "flowdir44.tif", "flowdir45.tif", "flowdir46.tif", "flowdir47.tif",
        "flowdir48.tif", "flowdir49.tif", "flowdir51.tif", "flowdir52.tif", "flowdir53.tif",
        "flowdir54.tif", "flowdir55.tif", "flowdir56.tif", "flowdir57.tif", "flowdir61.tif",
        "flowdir62.tif", "flowdir63.tif", "flowdir64.tif", "flowdir65.tif", "flowdir66.tif",
        "flowdir67.tif", "flowdir71.tif", "flowdir72.tif", "flowdir73.tif", "flowdir74.tif",
        "flowdir75.tif", "flowdir76.tif", "flowdir77.tif", "flowdir78.tif", "flowdir81.tif",
        "flowdir82.tif", "flowdir83.tif", "flowdir84.tif", "flowdir85.tif", "flowdir86.tif",
        "flowdir91.tif"
    ]
}

# 全局变量
progress_lock = threading.Lock()
global_progress = None

def setup_logger():
    """配置日志系统"""
    logger.remove()
    is_nohup = not sys.stdout.isatty()
    
    # 确保logs目录存在
    log_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "logs")
    os.makedirs(log_dir, exist_ok=True)
    log_file = os.path.join(log_dir, "mghydro_download.log")
    
    if is_nohup:
        # nohup环境：简洁日志
        logger.add(log_file, rotation="10 MB", retention="7 days",
                  format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}", level="INFO")
        logger.add(sys.stdout, format="{time:HH:mm:ss} | {message}", level="INFO", colorize=False)
    else:
        # 交互式环境：彩色日志
        logger.add(sys.stdout, format="<green>{time:HH:mm:ss}</green> | <level>{level}</level> | {message}", 
                  level="INFO", colorize=True)
        logger.add(log_file, rotation="10 MB", retention="7 days",
                  format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}", level="DEBUG")

def download_file(sub_path, filename):
    """下载单个文件"""
    global global_progress
    
    url = f"{BASE_URL}{sub_path}{filename}"
    local_path = os.path.join(BASE_DIR, sub_path, filename)
    
    # 创建目录
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    
    # 检查文件是否已存在（断点续传）
    start_byte = 0
    headers = {}
    mode = 'wb'
    
    if os.path.exists(local_path):
        start_byte = os.path.getsize(local_path)
        if start_byte > 0:
            headers['Range'] = f'bytes={start_byte}-'
            mode = 'ab'
            logger.debug(f"[{filename}] 续传，已有 {start_byte:,} 字节")
    
    try:
        response = requests.get(url, headers=headers, stream=True, timeout=30)
        
        if response.status_code == 404:
            logger.warning(f"[{filename}] 文件不存在")
            return
        
        if response.status_code == 416:  # 文件已完整
            logger.success(f"[{filename}] 已完整下载")
            return
        
        if response.status_code not in (200, 206):
            logger.error(f"[{filename}] HTTP {response.status_code}")
            return
        
        # 下载文件
        total_size = int(response.headers.get('content-length', 0)) + start_byte
        logger.info(f"[{filename}] 开始下载 ({total_size:,} 字节)")
        
        with open(local_path, mode) as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
        
        final_size = os.path.getsize(local_path)
        logger.success(f"[{filename}] 下载完成 ({final_size:,} 字节)")
        
    except Exception as e:
        logger.error(f"[{filename}] 下载失败: {e}")
    finally:
        with progress_lock:
            if global_progress:
                global_progress.update(1)

def main():
    """主函数"""
    global global_progress
    
    setup_logger()
    
    # 收集所有任务
    tasks = []
    for sub_path, files in FILES_TO_DOWNLOAD.items():
        for filename in files:
            tasks.append((sub_path, filename))
    
    logger.info(f"开始下载 {len(tasks)} 个文件")
    logger.info(f"使用 {MAX_WORKERS} 个并发线程")
    logger.info(f"目标目录: {BASE_DIR}")
    
    # 创建进度条
    global_progress = tqdm(total=len(tasks), desc="下载进度", unit="文件")
    
    try:
        with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            futures = [executor.submit(download_file, sub_path, filename) 
                      for sub_path, filename in tasks]
            
            # 等待完成
            success_count = 0
            for future in concurrent.futures.as_completed(futures):
                try:
                    future.result()
                    success_count += 1
                except Exception as e:
                    logger.error(f"任务异常: {e}")
    finally:
        if global_progress:
            global_progress.close()
    
    logger.success(f"下载完成！成功: {success_count}/{len(tasks)}")

if __name__ == "__main__":
    main()